{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "ename": "ValidationError",
     "evalue": "'{}' is not of type 'array'\n\nFailed validating 'type' in schema:\n    {'items': {'properties': {'data': {'items': {'type': 'number'},\n                                       'type': 'array'},\n                              'label': {'type': 'integer'},\n                              'type': {'type': 'string'}},\n               'type': 'object'},\n     'type': 'array'}\n\nOn instance:\n    '{}'",
     "output_type": "error",
     "traceback": [
      "\u001b[1;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[1;31mValidationError\u001b[0m                           Traceback (most recent call last)",
      "Cell \u001b[1;32mIn[12], line 24\u001b[0m\n\u001b[0;32m     16\u001b[0m input_labels \u001b[38;5;241m=\u001b[39m np\u001b[38;5;241m.\u001b[39marray([\u001b[38;5;241m1\u001b[39m, \u001b[38;5;241m1\u001b[39m, \u001b[38;5;241m2\u001b[39m])\n\u001b[0;32m     18\u001b[0m \u001b[38;5;66;03m# 使用 rembg 去除背景\u001b[39;00m\n\u001b[0;32m     19\u001b[0m \u001b[38;5;66;03m# output = remove(input)\u001b[39;00m\n\u001b[0;32m     20\u001b[0m \u001b[38;5;66;03m# output = remove(input, alpha_matting=True, alpha_matting_foreground_threshold=270,alpha_matting_background_threshold=20, alpha_matting_erode_size=11)\u001b[39;00m\n\u001b[0;32m     21\u001b[0m \u001b[38;5;66;03m# output = remove(input, only_mask=True)\u001b[39;00m\n\u001b[0;32m     22\u001b[0m \u001b[38;5;66;03m# output = remove(input, post_process_mask=True)\u001b[39;00m\n\u001b[0;32m     23\u001b[0m \u001b[38;5;66;03m# output = remove(input, session=session)\u001b[39;00m\n\u001b[1;32m---> 24\u001b[0m output \u001b[38;5;241m=\u001b[39m \u001b[43mremove\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;28;43minput\u001b[39;49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43msession\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43msession\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43minput_points\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43minput_points\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43minput_labels\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43minput_labels\u001b[49m\u001b[43m)\u001b[49m\n\u001b[0;32m     25\u001b[0m output\u001b[38;5;241m.\u001b[39msave(output_path)\n",
      "File \u001b[1;32mc:\\Users\\huigg\\anaconda3\\envs\\pytorch\\Lib\\site-packages\\rembg\\bg.py:257\u001b[0m, in \u001b[0;36mremove\u001b[1;34m(data, alpha_matting, alpha_matting_foreground_threshold, alpha_matting_background_threshold, alpha_matting_erode_size, session, only_mask, post_process_mask, bgcolor, *args, **kwargs)\u001b[0m\n\u001b[0;32m    254\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m session \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m:\n\u001b[0;32m    255\u001b[0m     session \u001b[38;5;241m=\u001b[39m new_session(\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mu2net\u001b[39m\u001b[38;5;124m\"\u001b[39m, \u001b[38;5;241m*\u001b[39margs, \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mkwargs)\n\u001b[1;32m--> 257\u001b[0m masks \u001b[38;5;241m=\u001b[39m \u001b[43msession\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mpredict\u001b[49m\u001b[43m(\u001b[49m\u001b[43mimg\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n\u001b[0;32m    258\u001b[0m cutouts \u001b[38;5;241m=\u001b[39m []\n\u001b[0;32m    260\u001b[0m \u001b[38;5;28;01mfor\u001b[39;00m mask \u001b[38;5;129;01min\u001b[39;00m masks:\n",
      "File \u001b[1;32mc:\\Users\\huigg\\anaconda3\\envs\\pytorch\\Lib\\site-packages\\rembg\\sessions\\sam.py:169\u001b[0m, in \u001b[0;36mSamSession.predict\u001b[1;34m(self, img, *args, **kwargs)\u001b[0m\n\u001b[0;32m    153\u001b[0m prompt \u001b[38;5;241m=\u001b[39m kwargs\u001b[38;5;241m.\u001b[39mget(\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124msam_prompt\u001b[39m\u001b[38;5;124m\"\u001b[39m, \u001b[38;5;124m\"\u001b[39m\u001b[38;5;132;01m{}\u001b[39;00m\u001b[38;5;124m\"\u001b[39m)\n\u001b[0;32m    154\u001b[0m schema \u001b[38;5;241m=\u001b[39m {\n\u001b[0;32m    155\u001b[0m     \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mtype\u001b[39m\u001b[38;5;124m\"\u001b[39m: \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124marray\u001b[39m\u001b[38;5;124m\"\u001b[39m,\n\u001b[0;32m    156\u001b[0m     \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mitems\u001b[39m\u001b[38;5;124m\"\u001b[39m: {\n\u001b[1;32m   (...)\u001b[0m\n\u001b[0;32m    166\u001b[0m     },\n\u001b[0;32m    167\u001b[0m }\n\u001b[1;32m--> 169\u001b[0m \u001b[43mvalidate\u001b[49m\u001b[43m(\u001b[49m\u001b[43minstance\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mprompt\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mschema\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mschema\u001b[49m\u001b[43m)\u001b[49m\n\u001b[0;32m    171\u001b[0m target_size \u001b[38;5;241m=\u001b[39m \u001b[38;5;241m1024\u001b[39m\n\u001b[0;32m    172\u001b[0m input_size \u001b[38;5;241m=\u001b[39m (\u001b[38;5;241m684\u001b[39m, \u001b[38;5;241m1024\u001b[39m)\n",
      "File \u001b[1;32mc:\\Users\\huigg\\anaconda3\\envs\\pytorch\\Lib\\site-packages\\jsonschema\\validators.py:1308\u001b[0m, in \u001b[0;36mvalidate\u001b[1;34m(instance, schema, cls, *args, **kwargs)\u001b[0m\n\u001b[0;32m   1306\u001b[0m error \u001b[38;5;241m=\u001b[39m exceptions\u001b[38;5;241m.\u001b[39mbest_match(validator\u001b[38;5;241m.\u001b[39miter_errors(instance))\n\u001b[0;32m   1307\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m error \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m:\n\u001b[1;32m-> 1308\u001b[0m     \u001b[38;5;28;01mraise\u001b[39;00m error\n",
      "\u001b[1;31mValidationError\u001b[0m: '{}' is not of type 'array'\n\nFailed validating 'type' in schema:\n    {'items': {'properties': {'data': {'items': {'type': 'number'},\n                                       'type': 'array'},\n                              'label': {'type': 'integer'},\n                              'type': {'type': 'string'}},\n               'type': 'object'},\n     'type': 'array'}\n\nOn instance:\n    '{}'"
     ]
    }
   ],
   "source": [
    "# 尝试去除背景\n",
    "import numpy as np\n",
    "from matplotlib import pyplot as plt\n",
    "import cv2\n",
    "from PIL import Image\n",
    "from rembg import remove, new_session\n",
    "\n",
    "input_path = '../masked_images/4.jpg'\n",
    "output_path = '../rembg/4.png'\n",
    "\n",
    "model_name = \"sam\"\n",
    "session = new_session(model_name)\n",
    "\n",
    "input = Image.open(input_path)\n",
    "input_points = np.array([[400, 350], [700, 400], [200, 400]])\n",
    "input_labels = np.array([1, 1, 2])\n",
    "\n",
    "# 使用 rembg 去除背景\n",
    "# output = remove(input)\n",
    "# output = remove(input, alpha_matting=True, alpha_matting_foreground_threshold=270,alpha_matting_background_threshold=20, alpha_matting_erode_size=11)\n",
    "# output = remove(input, only_mask=True)\n",
    "# output = remove(input, post_process_mask=True)\n",
    "# output = remove(input, session=session)\n",
    "output = remove(input, session=session, input_points=input_points, input_labels=input_labels)\n",
    "output.save(output_path)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "pytorch",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
